package com.blog.cmrpersonalblog.aspect;

import cn.dev33.satoken.stp.StpUtil;
import com.blog.cmrpersonalblog.annotation.RateLimit;
import com.blog.cmrpersonalblog.config.RateLimitConfig;
import com.blog.cmrpersonalblog.enums.LimitType;
import com.blog.cmrpersonalblog.exception.RateLimitException;
import jakarta.annotation.Resource;
import jakarta.servlet.http.HttpServletRequest;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import java.lang.reflect.Method;
import java.util.Collections;
import java.util.List;

/**
 * 限流切面
 * 使用Redis + Lua脚本实现滑动窗口限流算法
 */
@Slf4j
@Aspect
@Component
public class RateLimitAspect {
    
    @Resource
    private RedisTemplate<String, Object> redisTemplate;

    @Resource
    private RateLimitConfig rateLimitConfig;
    
    /**
     * Lua脚本：滑动窗口限流算法
     * 
     * KEYS[1]: 限流key
     * ARGV[1]: 时间窗口（秒）
     * ARGV[2]: 限流次数
     * ARGV[3]: 当前时间戳（毫秒）
     * 
     * 返回值：
     * 1: 允许访问
     * 0: 触发限流
     */
    private static final String LIMIT_SCRIPT = 
        "local key = KEYS[1]\n" +
        "local limit = tonumber(ARGV[1])\n" +
        "local window = tonumber(ARGV[2])\n" +
        "local current = tonumber(ARGV[3])\n" +
        "local expire_time = current - window * 1000\n" +
        "\n" +
        "-- 删除过期的记录\n" +
        "redis.call('ZREMRANGEBYSCORE', key, 0, expire_time)\n" +
        "\n" +
        "-- 获取当前窗口内的请求数\n" +
        "local current_count = redis.call('ZCARD', key)\n" +
        "\n" +
        "if current_count < limit then\n" +
        "    -- 未达到限流阈值，添加当前请求\n" +
        "    redis.call('ZADD', key, current, current)\n" +
        "    -- 设置过期时间\n" +
        "    redis.call('EXPIRE', key, window)\n" +
        "    return 1\n" +
        "else\n" +
        "    -- 达到限流阈值\n" +
        "    return 0\n" +
        "end";
    
    @Before("@annotation(rateLimit)")
    public void doBefore(JoinPoint joinPoint, RateLimit rateLimit) {
        // 检查是否启用限流
        if (!rateLimitConfig.getEnabled()) {
            return;
        }
        
        try {
            // 1. 构建限流key
            String key = buildKey(joinPoint, rateLimit);
            
            // 2. 获取限流参数
            int time = rateLimit.time();
            int count = rateLimit.count();
            
            // 3. 执行限流检查
            boolean allowed = checkRateLimit(key, count, time);
            
            if (!allowed) {
                // 触发限流
                String message = rateLimit.message();
                if (message == null || message.isEmpty()) {
                    message = String.format("访问过于频繁，请在%d秒后重试", time);
                }
                
                log.warn("触发限流: key={}, limit={}/{} seconds", key, count, time);
                throw new RateLimitException(message, time, count, time);
            }
            
        } catch (RateLimitException e) {
            throw e;
        } catch (Exception e) {
            log.error("限流检查失败", e);
            // 限流检查失败时，默认放行（避免影响正常业务）
        }
    }
    
    /**
     * 构建限流key
     */
    private String buildKey(JoinPoint joinPoint, RateLimit rateLimit) {
        StringBuilder key = new StringBuilder(rateLimitConfig.getKeyPrefix());
        
        // 添加方法标识
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        String className = method.getDeclaringClass().getSimpleName();
        String methodName = method.getName();
        
        if (rateLimit.key() != null && !rateLimit.key().isEmpty()) {
            key.append(rateLimit.key());
        } else {
            key.append(className).append(":").append(methodName);
        }
        
        // 根据限流类型添加后缀
        LimitType limitType = rateLimit.limitType();
        switch (limitType) {
            case IP:
                key.append(":").append(getIpAddress());
                break;
            case USER:
                try {
                    Object loginId = StpUtil.getLoginIdDefaultNull();
                    if (loginId != null) {
                        key.append(":user:").append(loginId);
                    } else {
                        // 未登录用户按IP限流
                        key.append(":ip:").append(getIpAddress());
                    }
                } catch (Exception e) {
                    // 获取用户ID失败，按IP限流
                    key.append(":ip:").append(getIpAddress());
                }
                break;
            case GLOBAL:
                // 全局限流，不添加后缀
                break;
        }
        
        return key.toString();
    }
    
    /**
     * 执行限流检查
     * 
     * @param key 限流key
     * @param count 限流次数
     * @param time 时间窗口（秒）
     * @return true-允许访问，false-触发限流
     */
    private boolean checkRateLimit(String key, int count, int time) {
        try {
            DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
            redisScript.setScriptText(LIMIT_SCRIPT);
            redisScript.setResultType(Long.class);
            
            List<String> keys = Collections.singletonList(key);
            Long result = redisTemplate.execute(
                redisScript,
                keys,
                count,
                time,
                System.currentTimeMillis()
            );
            
            return result != null && result == 1L;
        } catch (Exception e) {
            log.error("执行限流脚本失败: key={}", key, e);
            // 脚本执行失败，默认放行
            return true;
        }
    }
    
    /**
     * 获取客户端IP地址
     */
    private String getIpAddress() {
        try {
            ServletRequestAttributes attributes = 
                (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
            if (attributes == null) {
                return "unknown";
            }
            
            HttpServletRequest request = attributes.getRequest();
            String ip = request.getHeader("X-Forwarded-For");
            
            if (ip == null || ip.isEmpty() || "unknown".equalsIgnoreCase(ip)) {
                ip = request.getHeader("X-Real-IP");
            }
            if (ip == null || ip.isEmpty() || "unknown".equalsIgnoreCase(ip)) {
                ip = request.getHeader("Proxy-Client-IP");
            }
            if (ip == null || ip.isEmpty() || "unknown".equalsIgnoreCase(ip)) {
                ip = request.getHeader("WL-Proxy-Client-IP");
            }
            if (ip == null || ip.isEmpty() || "unknown".equalsIgnoreCase(ip)) {
                ip = request.getRemoteAddr();
            }
            
            // 处理多个IP的情况（取第一个）
            if (ip != null && ip.contains(",")) {
                ip = ip.split(",")[0].trim();
            }
            
            return ip;
        } catch (Exception e) {
            log.error("获取IP地址失败", e);
            return "unknown";
        }
    }
}

